"""
Visualization utilities: price + signals, MACD, equity curve, drawdown
"""
from typing import Optional
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import numpy as np


def setup_chinese():
    plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'Songti SC', 'Heiti TC']
    plt.rcParams['axes.unicode_minus'] = False


def plot_results(df: pd.DataFrame, eq_df: pd.DataFrame, trades_df: pd.DataFrame, savepath: Optional[str] = None):
    setup_chinese()
    df = df.copy()
    df['日期'] = pd.to_datetime(df['日期'])
    df = df.sort_values('日期').reset_index(drop=True)
    fig, axes = plt.subplots(4, 1, figsize=(14, 12), sharex=True, gridspec_kw={'height_ratios':[2,1,1,1]})

    # price
    ax = axes[0]
    ax.plot(df['日期'], df['收盘'], label='收盘')
    # mark trades
    if not trades_df.empty:
        for _, t in trades_df.iterrows():
            if t['type'] == 'entry':
                ax.annotate('买', xy=(t['date'], t['price']), xytext=(0, 10), textcoords='offset points', arrowprops=dict(arrowstyle='->', color='g'))
            else:
                ax.annotate('卖', xy=(t['date'], t['price']), xytext=(0, -15), textcoords='offset points', arrowprops=dict(arrowstyle='->', color='r'))
    ax.set_title('价格与买卖信号')
    ax.legend()

    # MACD
    ax2 = axes[1]
    ax2.plot(df['日期'], df['macd_line'], label='MACD线')
    ax2.plot(df['日期'], df['macd_signal'], label='信号线')
    ax2.bar(df['日期'], df['macd_hist'], label='柱状图')
    ax2.set_title('MACD')
    ax2.legend()

    # RSI
    ax3 = axes[2]
    ax3.plot(df['日期'], df['rsi'], label='RSI')
    ax3.axhline(70, color='r', linestyle='--')
    ax3.axhline(30, color='g', linestyle='--')
    ax3.set_title('RSI')
    ax3.legend()

    # equity
    ax4 = axes[3]
    ax4.plot(eq_df['date'], eq_df['equity'], label='净值')
    ax4.set_title('资金净值曲线')
    ax4.legend()

    plt.tight_layout()
    if savepath:
        fig.savefig(savepath)
    return fig
